import multiprocessing
import time

import wget
import os
from gxl_ai_utils.utils import utils_file

def download_audio_with_wget(audio_url, save_dir):
    """
    使用 python-wget 下载音频并指定保存目录
    Args:
        audio_url: 音频文件的 HTTP 链接
        save_dir: 保存音频的目标目录
    """
    # 1. 处理保存目录：不存在则创建
    os.makedirs(save_dir, exist_ok=True)

    # 2. 提取音频文件名，拼接保存路径
    audio_filename = audio_url.split("/")[-1]
    save_path = os.path.join(save_dir, audio_filename)

    # 3. 调用 wget.download 下载（out 参数指定保存路径）
    try:
        wget.download(audio_url, out=save_path)
        return save_path
    except Exception as e:
        print(f"下载失败！原因：{str(e)}")
        raise e


# ------------------- 调用示例 -------------------
audio_url = "http://9.223.97.217/dat_en_bpo/librispeech/data/test-clean/1089/134686/1089-134686-0009.opus"
save_directory = "./data/audio_downloads_wget"  # 自定义保存目录
utils_file.makedir_sil(save_directory)
# 执行下载
download_audio_with_wget(audio_url, save_directory)

root_dir = "/apdcephfs_qy3/share_976139/users/xuelonggeng/data/asr_test"
test_data_list=[
   "/apdcephfs_qy3/share_976139/users/xuelonggeng/data/asr_test/librispeechh_clean/test_20220407_librispeech_clean_org_fbankhires.jsonl",
   # "/teaspeech_ceph/share_976139/users/kaixunhuang/buf/test_dir/test_20220407_librispeech_other_org_fbankhires_local.jsonl", # 本地路径
   "/teaspeech_ceph/share_976139/users/kaixunhuang/buf/test_dir/test_20211229_wenetspeech_meeting_org_fbankhires.jsonl",
   "/teaspeech_ceph/share_976139/users/kaixunhuang/buf/test_dir/test_20241119_aishell2_android.jsonl",
   "/teaspeech_ceph/share_976139/users/kaixunhuang/buf/test_dir/test_20241119_aishell2_ios.jsonl",
   "/teaspeech_ceph/share_976139/users/kaixunhuang/buf/test_dir/test_20241119_aishell2_mic.jsonl",
    # "/apdcephfs_qy3/share_976139/users/kaixunhuang/buf/test_dir/test_20250429_wenetspeech_test_net.jsonl"
    # "/apdcephfs_qy3/share_976139/users/kaixunhuang/buf/test_dir/test_aishell4.jsonl"
    # "/apdcephfs_qy3/share_976139/users/kaixunhuang/buf/test_dir/test_en_gigaspeech.jsonl"
    # "/apdcephfs_qy3/share_976139/users/xuelonggeng/data/asr_test/aishell_test/data.list"
]

testset_name_list=[
   "librispeech_clean",
   # "librispeech_other",
   "wenetspeech_meeting",
   "aishell2_android",
   "aishell2_ios",
   "aishell2_mic",
    # "testnet_test"
    # "aishell4_test"
    # "gigaspeech_test"
    # "aishell1_test"
]
from functools import partial
class GXLMultiprocessingWithReturn:
    def __init__(self, num_processes):
        self.num_processes = num_processes
        self.ctx = multiprocessing.get_context('spawn')

    def run(self, func, big_dict_or_list, **kwargs):
        assert big_dict_or_list is not None, "big_dict_or_list 不能为空！"
        assert callable(func), "func 必须是一个可调用对象！"
        assert isinstance(big_dict_or_list, (dict, list)), "big_dict_or_list 必须是一个字典或列表！"
        if isinstance(big_dict_or_list, dict):
            split_list =utils_file.do_split_dict(big_dict_or_list, self.num_processes)
        else:
            split_list =utils_file.do_split_list(big_dict_or_list, self.num_processes)
        utils_file.logging_info(f'开始多进程处理，共有 {len(split_list)} 个子任务。')
        time_start = time.time()
        wrapped_func = partial(func, **kwargs)
        results_from_all_processes = []
        with multiprocessing.Pool(self.num_processes) as pool:
            results_from_all_processes = pool.map(wrapped_func, split_list)
        end_time = time.time()
        print(f"--- 所有进程处理完毕，耗时: {end_time - time_start:.2f} 秒 ---")
        return results_from_all_processes

def little_func(little_dict_list, save_dir):
    """"""
    new_dict_list = []
    for little_dict in utils_file.tqdm(little_dict_list, desc="下载音频", total=len(little_dict_list)):
        wav_path = little_dict['wav']
        try:
            new_wav_path = download_audio_with_wget(wav_path, save_dir)
            little_dict['wav'] = new_wav_path
            new_dict_list.append(little_dict)
        except Exception as e:
            print(f"下载失败！原因：{str(e)}")
            continue
    return new_dict_list

for test_data, testset_name in zip(test_data_list, testset_name_list):
    save_dir = os.path.join(root_dir, testset_name, "wavs")
    utils_file.logging_info(f"开始处理 {testset_name} 数据集, save_dir: {save_dir}")
    utils_file.makedir_sil(save_dir)
    dict_list = utils_file.load_dict_list_from_jsonl(test_data)
    num_thread = 50
    # runner = GXLMultiprocessingWithReturn(num_thread)
    # res = runner.run(little_func, dict_list, save_dir=save_dir)
    # total_new_dict_list = []
    # for res_i in res:
    #     total_new_dict_list.extend(res_i)
    old_path = os.path.join(save_dir, testset_name, 'data.jsonl')
    output_path_jsonl_path = os.path.join(root_dir, testset_name, 'data.jsonl')
    # utils_file.write_dict_list_to_jsonl(total_new_dict_list, output_path_jsonl_path)
    utils_file.copy_file(old_path, output_path_jsonl_path)

